from .module import DATASET
from .base_dataset import BaseDataset
from progressbar import *
import cv2
import os
import numpy as np

@DATASET.register_module
class CUB(BaseDataset):

    def __init__(self, path, preprocess_pipe, mode='train'):
        super(CUB, self).__init__(path, preprocess_pipe, mode)

    def load_data(self):
        images_path = '{}/images.txt'.format(self.data_path)

        image_ids = []
        with open(images_path, 'r') as f:
            line = f.readline()
            widgets = ['CUB loading: ', Percentage(), ' ', Bar('|'), ' ', Timer(),
                       ' ', ETA(), ' ']
            pbar = ProgressBar(widgets=widgets, maxval=12000).start()
            count = 0
            while line != '':
            # for _ in range(1000):
                line = line.strip('\n')
                id, img_path = line.split(' ')
                img = cv2.imread(os.path.join(self.data_path, 'images', img_path))
                #to rgb
                r = img[:, :, 2]
                img[:, :, 2] = img[:, :, 0]
                img[:, :, 0] = r

                self.input_data.append(img)
                image_ids.append(id)
                line = f.readline()
                count += 1
                pbar.update(count)
            pbar.finish()
            f.close()

        class_label_path = '{}/image_class_labels.txt'.format(self.data_path)
        cls_dict = {}
        with open(class_label_path, 'r') as f:
            line = f.readline()
            while line != '':
                line = line.strip('\n')
                id, cls = line.split(' ')
                line = f.readline()
                cls_dict[id] = int(cls)
        for id in image_ids:
            self.annotation.append(cls_dict[id] - 1)

        train_test_split_path = '{}/train_test_split.txt'.format(self.data_path)
        mode_dict = {}
        with open(train_test_split_path, 'r') as f:
            line = f.readline()
            while line != '':
                line = line.strip('\n')
                id, mode = line.split(' ')
                line = f.readline()
                mode_dict[id] = mode
        for i, id in enumerate(image_ids):
            if mode_dict[id] == '1':
                self.train_indexes.append(i)
            else:
                self.test_indexes.append(i)

        bounding_boxes_path = '{}/bounding_boxes.txt'.format(self.data_path)
        bounding_boxes_dict = {}
        with open(bounding_boxes_path, 'r') as f:
            line = f.readline()
            while line != '':
                line = line.strip('\n')
                id, x, y, w, h = line.split(' ')
                x = int(float(x))
                y = int(float(y))
                w = int(float(w))
                h = int(float(h))
                line = f.readline()
                bounding_boxes_dict[id] = [x, y, x+w, y+h]
        for i, id in enumerate(image_ids):
            self.annotation[i] = np.array([self.annotation[i]] + bounding_boxes_dict[id], dtype=np.int32)

        print('CUB:{:d} samples for training, {:d} samples for test'
              .format(len(self.train_indexes), len(self.test_indexes)))